#include "rsa.h"
#include "kprint.h"
#define MOD( a, b ) ((a)%(b))

void * __gxx_personality_v0=0;
void * _Unwind_Resume =0;

char buffer[ 1024 ];
const int MAX_DIGITS = 50;
int i, j = 0;

void rsa_printf(const char *str)
{
    // printf( "%s", str );
    kprintf(str);
}

// This should totally be in the math library.
myuint64_t gcd( myuint64_t a, myuint64_t b )
{
    myuint64_t c;
    while ( a != 0 )
    {
        c = a;
        a = MOD( b, a );
        b = c;
    }
    return b;
}

myuint64_t ExtEuclid( myuint64_t a, myuint64_t b )
{
    myuint64_t x = 0, y = 1, u = 1, v = 0, gcd = b, m, n, q, r;
    while ( a != 0 )
    {
        q = gcd / a;
        r = MOD( gcd, a );
        m = x - u * q;
        n = y - v * q;
        gcd = a;
        a = r;
        x = u;
        y = v;
        u = m;
        v = n;
    }
    return y;
}

static inline myuint64_t modmult( myuint64_t a, myuint64_t b, myuint64_t mod )
{
    // this is necessary since we will be dividing by a
    if ( a == 0 )
    {
        return 0;
    }
    myuint64_t product = a * b;
    // if multiplication does not overflow, we can use it
    if ( product / a == b )
    {
        return MOD( product, mod );
    }
    // if a % 2 == 1 i. e. a >> 1 is not a / 2
    if ( (a & 1) != 0 )
    {
        product = modmult( ( a >> 1 ), b, mod );
        if ( ( product << 1 ) > product )
        {
            product <<= 1;
            return MOD( ( MOD( product, mod ) + b ), mod );
            // return ( ( ( product << 1 ) % mod ) + b ) % mod;
        }
    }
    // implicit else
    product = modmult( ( a >> 1 ), b, mod );
    if ( ( product << 1 ) > product )
    {
        return MOD( ( product << 1 ), mod );
        // return ( product << 1 ) % mod;
    }
    // implicit else: this is about 10x slower than the code above, but it will
    // not overflow
    myuint64_t sum;
    sum = 0;
    while ( b > 0 )
    {
        if ( (b & 1 ) != 0 )
            sum = MOD( ( sum + a ), mod );
        a = MOD( a >> 1, mod );
        b >>= 1;
    }
    return sum;
}

myuint64_t rsa_modExp( myuint64_t b, myuint64_t e, myuint64_t m )
{
    myuint64_t product;
    product = 1;
    if ( b < 0 || e < 0 || m <= (myuint64_t)0 )
    {
        return -1;
    }
    b = MOD( b, m );
    while ( e > 0 )
    {
        if ( (e & 1 )!= 0)
        {
            product = modmult( product, b, m );
        }
        b = modmult( b, b, m );
        e >>= 1;
    }
    return product;
}

void rsa_gen_keys( struct public_key_class *pub, private_key_class *priv )
{
    myuint64_t p = 0;
    myuint64_t q = 0;

    // values of e should be sufficiently large to protect against naive attacks
    myuint64_t e = ( 2 << 16 ) + 1;
    myuint64_t d = 0;
    myuint64_t max = 0;
    myuint64_t phi_max = 0;

    p = 59921;
    q = 34631;
    max = p * q;
    phi_max = ( p - 1 ) * ( q - 1 );
    if ( !( p && q ) || ( p == q ) || ( gcd( phi_max, e ) != 1 ) )
    {
        rsa_printf( "error" );
    }

    d = ExtEuclid( phi_max, e );
    // while ( d < 0 )//由于使用的是uint64_t，因此修改如下(直接对d进行强制赋值)
    // {
    //     d = d + phi_max;
    // }
    d = 1860581537;

    // rsa_printf("primes are %lld and %lld\n",(uint64_t)p, (uint64_t )q);
    // We now store the public / private keys in the appropriate structs
    pub->modulus = max;
    pub->exponent = e;

    priv->modulus = max;
    priv->exponent = d;
}

bool rsa_encrypt( const uint8_t *message, const size_t message_size,
                  const struct public_key_class *pub, myuint64_t *encrypted )
{
    if ( encrypted == NULL )
    {
        char text[] = { "Error: Heap allocation failed.\n" };
        rsa_printf(text);
        return false;
    }
    for (size_t i = 0; i < message_size; i++ )
    {
        if ( ( encrypted[ i ] = rsa_modExp( message[ i ], pub->exponent,
                                            pub->modulus ) ) == -1 )
        {
            char text[] = { "Error: rsa_modExp"};
            rsa_printf(text);
            return false;
        }
    }
    return true;
}

bool rsa_decrypt( const myuint64_t *message, const size_t message_size,
                  const private_key_class *priv, uint8_t *decrypted )
{
    if ( MOD( message_size, sizeof( myuint64_t ) ) != 0 )
    {
        char text[] = {"Error: message_size is not divisible by sizeof( uint64_t ), so "
            "cannot be "
            "output of rsa_encrypt\n" };
        rsa_printf(text);
        return false;
    }
    // We allocate space to do the decryption (temp) and space for the output as
    // a char array (decrypted)
    // char *decrypted = malloc( message_size / sizeof( uint64_t ) );
    char temp[ message_size ];
    if ( ( decrypted == NULL ) || ( temp == NULL ) )
    {
        char text[] = { "Error: Heap allocation failed.\n" };
        rsa_printf( text );
        return false;
    }
    // Now we go through each 8-byte chunk and decrypt it.
    for (size_t i = 0; i < message_size / sizeof(myuint64_t); i++ )
    {
        if ( ( temp[ i ] = rsa_modExp( message[ i ], priv->exponent,
                                       priv->modulus ) ) == -1 )
        {
            return false;
        }
    }
    // The result should be a number in the char range, which gives back the
    // original byte. We put that into decrypted, then return.
    for (size_t i = 0; i < message_size / 8; i++ )
    {
        decrypted[ i ] = temp[ i ];
    }
    return true;
}
